// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/varint.h>
#include <melon/init/init.h>
#include <pollux/common/memory/memory.h>
#include <pollux/dwio/common/data_buffer_holder.h>
#include <pollux/dwio/common/range.h>
#include <pollux/dwio/dwrf/common/encoder_util.h>
#include <pollux/dwio/dwrf/common/int_encoder.h>

using namespace kumo::pollux::dwio::common;
using namespace kumo::pollux;
using namespace kumo::pollux::dwrf;
using namespace kumo::pollux::memory;

static size_t generateAutoId(int64_t startId, int64_t count) {
  size_t capacity = count * melon::kMaxVarintLength64;
  auto pool = memory::memoryManager()->addLeafPool();
  DataBufferHolder holder{*pool, capacity};
  auto output = std::make_unique<BufferedOutputStream>(holder);
  auto encoder =
      createDirectEncoder<true>(std::move(output), true, sizeof(int64_t));

  for (int64_t i = 0; i < count; i++) {
    encoder->writeValue(startId + i);
  }
  return encoder->flush();
}

static size_t generateAutoId2(int64_t startId, int64_t count) {
  size_t capacity = count * melon::kMaxVarintLength64;
  auto pool = memory::memoryManager()->addLeafPool();
  DataBufferHolder holder{*pool, capacity};
  auto output = std::make_unique<BufferedOutputStream>(holder);
  auto encoder =
      createDirectEncoder<true>(std::move(output), true, sizeof(int64_t));

  int64_t buffer[1024];
  int64_t currentId = startId;
  int64_t countRemaining = count;
  while (countRemaining > 0) {
    int64_t bufCount = std::min(countRemaining, (int64_t)1024);
    for (int64_t i = 0; i < bufCount; ++i) {
      buffer[i] = currentId++;
    }
    encoder->add(buffer, common::Ranges::of(0, bufCount), nullptr);
    countRemaining -= bufCount;
  }
  return encoder->flush();
}

MELON_ALWAYS_INLINE static int32_t findSetBitsOld(uint64_t value) {
  if (value < (1ul << 14)) {
    if (value < (1ul << 7)) {
      return 1;
    }
    return 2;
  } else if (value < (1ul << 42)) {
    if (value < (1ul << 21)) {
      return 3;
    } else if (value < (1ul << 28)) {
      return 4;
    } else if (value < (1ul << 35)) {
      return 5;
    }
    return 6;
  } else {
    if (value < (1ul << 49)) {
      return 7;
    } else if (value < (1ul << 56)) {
      return 8;
    } else if (value < (1ul << 63)) {
      return 9;
    }
    return 10;
  }
}

MELON_ALWAYS_INLINE static int32_t findSetBitsNew(uint64_t value) {
  int32_t leadingZeros = __builtin_clzll(value | 1);
  DKCHECK(leadingZeros <= 63);
  // bytes in varint can be calculated as (70 - leadingZeros)/7;
  switch (leadingZeros) {
    case 0:
      return 10;
    case 1 ... 7:
      return 9;
    case 8 ... 14:
      return 8;
    case 15 ... 21:
      return 7;
    case 22 ... 28:
      return 6;
    case 29 ... 35:
      return 5;
    case 36 ... 42:
      return 4;
    case 43 ... 49:
      return 3;
    case 50 ... 56:
      return 2;
    case 57 ... 63:
      return 1;
  }
  DWIO_RAISE(melon::sformat(
      "Unexpected leading zeros {} for value {}", leadingZeros, value));
}

size_t iters = 2000;

BENCHMARK(findSetBitsOld) {
  uint64_t value = 0;
  for (int64_t i = 0; i < iters; i++) {
    value += INT16_MAX;
    auto result = findSetBitsOld(value);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK_RELATIVE(findSetBitsNew) {
  uint64_t value = 0;
  for (int64_t i = 0; i < iters; i++) {
    value += INT16_MAX;
    auto result = findSetBitsNew(value);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK(findSetBitsOld_low) {
  uint64_t value = 0;
  for (int64_t i = 0; i < iters; i++) {
    value++;
    auto result = findSetBitsOld(value);
    melon::doNotOptimizeAway(result);
    value = value & 0xff;
  }
}

BENCHMARK_RELATIVE(findSetBitsNew_low) {
  uint64_t value = 0;
  for (int64_t i = 0; i < iters; i++) {
    value++;
    auto result = findSetBitsNew(value);
    melon::doNotOptimizeAway(result);
    value = value & 0xff;
  }
}

BENCHMARK(GenerateAutoIdOld_0) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId(0, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK_RELATIVE(GenerateAutoIdNew_0) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId2(0, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK(GenerateAutoIdOld_32) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId(INT32_MAX, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK_RELATIVE(GenerateAutoIdNew_32) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId2(INT32_MAX, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK(GenerateAutoIdOld_64) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId(INT64_MAX - 500'000, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

BENCHMARK_RELATIVE(GenerateAutoIdNew_64) {
  for (int64_t i = 0; i < iters; i++) {
    auto result = generateAutoId2(INT64_MAX - 500'000, 100'000);
    melon::doNotOptimizeAway(result);
  }
}

int32_t main(int32_t argc, char* argv[]) {
  melon::Init init{&argc, &argv};
  memory::MemoryManager::initialize({});
  melon::runBenchmarks();
  return 0;
}
